﻿namespace Cellfish.Toolbox.Tasks
{
    using System;
    using System.Threading;
    using System.Threading.Tasks;

    /// <summary>
    /// Extension methods to make tasks "cancellable".
    /// </summary>
    public static class WithCancellationExtensions
    {

        /// <summary>
        /// Adds cancellation possibility to a task.
        /// </summary>
        /// <remarks>Note that the task is not stopped when cancelled.</remarks>
        /// <param name="task">The task that should have cancellation added.</param>
        /// <param name="cancellationToken">The cancellation token.</param>
        /// <returns>A task that represents the completion of the supplied task.</returns>
        /// <exception cref="System.ArgumentNullException">If task is null.</exception>
        public static Task WithCancellation(this Task task, CancellationToken cancellationToken)
        {
            if (task == null)
            {
                throw new ArgumentNullException("task");
            }

            if (task.IsCompleted || !cancellationToken.CanBeCanceled)
            {
                return task;
            }

            return AddCancellation(task.WithResult(true), cancellationToken);
        }

        /// <summary>
        /// Adds cancellation possibility to a task.
        /// </summary>
        /// <remarks>Note that the task is not stopped when cancelled.</remarks>
        /// <typeparam name="T">Result type of the task.</typeparam>
        /// <param name="task">The task that should have cancellation added.</param>
        /// <param name="cancellationToken">The cancellation token.</param>
        /// <returns>A task that represents the completion of the supplied task.</returns>
        /// <exception cref="System.ArgumentNullException">If task is null.</exception>
        public static Task<T> WithCancellation<T>(this Task<T> task, CancellationToken cancellationToken)
        {
            if (task == null)
            {
                throw new ArgumentNullException("task");
            }

            if (task.IsCompleted || !cancellationToken.CanBeCanceled)
            {
                return task;
            }

            return AddCancellation(task, cancellationToken);
        }

        private static Task<T> AddCancellation<T>(Task<T> task, CancellationToken cancellationToken)
        {
            var taskCompletionSource = new TaskCompletionSource<T>();
            var cancellationRegistration = cancellationToken.Register(() => taskCompletionSource.TrySetCanceled());
            task.ContinueWith(
                t =>
                {
                    cancellationRegistration.Dispose();
                    if (t.IsCanceled)
                    {
                        taskCompletionSource.TrySetCanceled();
                    }
                    else if (t.IsFaulted)
                    {
                        taskCompletionSource.TrySetException(t.Exception.InnerException);
                    }
                    else
                    {
                        taskCompletionSource.TrySetResult(t.Result);
                    }
                }, 
                TaskContinuationOptions.ExecuteSynchronously);
            return taskCompletionSource.Task;
        }
    }
}
